from enum import IntEnum

from core.aparser import Parser
from core.ascanner import Scanner
from core.bytecodes import LOAD_CONST, ADD, SUB, MUL, SET_GLOBAL, LOAD_GLOBAL, PRINT, CALL, RETURN, disassemble_codes, \
    SET_LOCAL, LOAD_LOCAL
from core.exprs import Literal, Application, Operator, Variable, DefineVar, Body, DefineFunction
from core.mtoken import TokenType
from core.objects import ExecutorObject, AS_STR, AS_NUM, AS_BOOL, Function, Value

scnaner = Scanner()
parser = Parser()


class ScopeCode:
    def __init__(self):
        self.codes = []


class VisitorScope(IntEnum):
    NONE = 0
    DEFINE_VAR = 1
    DEFINE_FUNC = 2


class TransContext:
    def __init__(self, parent: 'TransContext'):
        self.parent = parent
        self.code_frame: list[ScopeCode] = [ScopeCode()]
        self.globals = {}
        self.locals = {}
        self.functions = {}
        self.constants: list[Value] = []
        self.status = None
        self.deep = 0

        if parent is not None:
            self.globals = parent.globals
            self.functions = parent.functions
            self.constants = parent.constants
            self.code_frame = parent.code_frame
            self.code_frame.append(ScopeCode())
            self.deep = parent.deep + 1

    def curr_code_frame(self) -> ScopeCode:
        return self.code_frame[-1]

    def pop_code_frame(self) -> ScopeCode:
        return self.code_frame.pop()

    def add_constant(self, value) -> int:
        self.constants.append(value)
        return len(self.constants) - 1

    def search_local(self, var_name):
        if var_name in self.locals:
            return True
        else:
            if self.parent is not None:
                return self.parent.search_local(var_name)
            else:
                return False


class BytecodeTranslatorVisitor:

    def __init__(self, ctx: TransContext):
        self.ctx = ctx
        self.deep = 0
        self.constant_map = {}

    def visit(self, expr):
        if isinstance(expr, Literal):
            if isinstance(expr.value, str):
                constant = AS_STR(expr.value)
            elif isinstance(expr.value, int) or isinstance(expr.value, float):
                constant = AS_NUM(expr.value)
            elif isinstance(expr.value, bool):
                constant = AS_BOOL(expr.value)
            else:
                raise Exception("not support type")
            index = self.ctx.add_constant(constant)
            self.emit_constant(index)
        elif isinstance(expr, Variable):
            if (expr.name.value == "print"):
                self.emit_byte(PRINT)
                return
            var_name = AS_STR(expr.name.value)
            index = self.ctx.add_constant(var_name)
            self.emit_constant(index)
            if self.ctx.status != "define":
                if self.ctx.search_local(var_name):
                    self.emit_byte(LOAD_LOCAL)
                else:
                    self.emit_byte(LOAD_GLOBAL)
        elif isinstance(expr, Application):
            for arg in expr.sequences:
                self.visit(arg)

            self.visit(expr.first)
            if isinstance(expr.first, Variable) and expr.first.name.value != "print" and self.ctx != "define":
                self.emit_byte(CALL)
        elif isinstance(expr, Operator):
            if expr.operator.ttype == TokenType.PLUS:
                self.emit_byte(ADD)
            elif expr.operator.ttype == TokenType.MINUS:
                self.emit_byte(SUB)
            elif expr.operator.ttype == TokenType.STAR:
                self.emit_byte(MUL)
        elif isinstance(expr, DefineVar):
            self.ctx.status = "define"
            # 添加变量名作为常量
            var_str = AS_STR(expr.name.value)
            index = self.ctx.add_constant(var_str)
            self.emit_constant(index)
            # 生成复制的表达式字节码
            self.visit(expr.value)
            if self.ctx.deep > 0:
                self.emit_byte(SET_LOCAL)
                self.ctx.locals[var_str] = index
            else:
                self.emit_byte(SET_GLOBAL)
            self.ctx.status = None
        elif isinstance(expr, Body):
            for define in expr.defines:
                self.visit(define)
            for expr in expr.expressions:
                self.visit(expr)
        elif isinstance(expr, DefineFunction):
            func_name_str = AS_STR(expr.name.value)
            self.push_ctx()
            args = []
            for argument in expr.arguments:
                param_str = AS_STR(argument.name.value)
                args.append(param_str)
                self.ctx.locals[param_str] = 1
            self.visit(expr.body)
            self.emit_byte(RETURN)


            self.pop_ctx()


            code_frame = self.ctx.pop_code_frame()
            procession = Function(func_name_str, args, code_frame.codes)
            self.add_global(func_name_str, procession)
        return

    def emit_constant(self, index):
        self.emit_bytes(LOAD_CONST, index)

    def emit_byte(self, byte):
        code_scope = self.ctx.curr_code_frame()
        code_scope.codes.append(byte)

    def emit_bytes(self, byte1, byte2):
        self.emit_byte(byte1)
        self.emit_byte(byte2)

    def emit_global(self):
        self.emit_byte(SET_GLOBAL)

    def add_global(self, func_name_str, procession):
        self.ctx.globals[func_name_str] = procession

    def push_ctx(self):
        sub_ctx = TransContext(self.ctx)
        self.ctx = sub_ctx

    def pop_ctx(self):
        self.ctx = self.ctx.parent


def build_executor_obj(ctx: TransContext):
    obj = ExecutorObject()
    obj.constants = ctx.constants
    obj.globals = ctx.globals
    obj.functions = ctx.functions
    obj.codes = ctx.curr_code_frame().codes
    return obj


class Compiler:
    def reset(self):
        self.trans_ctx = TransContext(None)
        self.visitor = BytecodeTranslatorVisitor(self.trans_ctx)

    def compile(self, code: str) -> ExecutorObject:
        self.reset()
        tokens = scnaner.scan_tokens(code)
        exprs = parser.parse(tokens)
        for expr in exprs:
            self.visit(expr)

        return build_executor_obj(self.trans_ctx)

    def visit(self, expr):
        self.visitor.visit(expr)


if __name__ == '__main__':
    compiler = Compiler()
    # code = """
    #     (define (add x y) (+ x y))
    #     (print (add 2 3))
    # """
    code = """
        (define y 2)
        (define (func1 x)
            (+ x y)
        )
    """
    obj = compiler.compile(code)
    print("function main")
    disassemble_codes(obj.codes)

    for key, value in obj.globals.items():
        print(f"function {key.val}")
        disassemble_codes(value.body)